/*
This file is part of MMM.

MMM is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

MMM is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with MMM.  If not, see <http://www.gnu.org/licenses/>.
*
* @package    MMMTools
* @author     Christian Mandery
* @copyright  2015 High Performance Humanoid Technologies (H2T), Karlsruhe, Germany
*
*/

#include <MMM/MathTools.h>
#include <MMMSimoxTools/MMMSimoxTools.h>
#include <MMM/Motion/Plugin/KinematicPlugin/KinematicSensor.h>
#include <MMM/Motion/Plugin/ModelPosePlugin/ModelPoseSensor.h>
#include <MMM/Motion/Plugin/MoCapMarkerPlugin/MoCapMarkerSensor.h>
#include <MMM/Exceptions.h>

#include "FramewiseLocalStrategy.h"
#include "MinimumJerkStrategy.h"

#include <map>

#include "NloptMotionConverter.h"

using namespace MMM;

NloptMotionConverter::NloptMotionConverter(MotionPtr inputMotion, ModelPtr outputModel, ModelPtr outputModelUnprocessed, ModelProcessorPtr
                                           outputModelProcessor, const std::string &configFile, const std::string &convertedMotionName,
                                           OptimizationStrategy optimizationStrategy)
    :   MotionConverter(inputMotion, outputModel, outputModelUnprocessed, outputModelProcessor, configFile, convertedMotionName),
        optimizationStrategy(optimizationStrategy),
        markerScaleFactor(1.0f),
        targetModelSize(1.0f)
{
    if (!inputMotion || !outputModel) throw Exception::MMMException("Input motion and output model should not be null");

    try {
        readConverterConfig();
    } catch (Exception::MMMException &e) {
        throw Exception::MMMException("Exception while reading nlopt converter config file! " + std::string(e.what()));
    }
}

void NloptMotionConverter::readConverterConfig() {
    MMM_INFO << "Reading XML converter config." << std::endl;

    RapidXMLReaderPtr reader = RapidXMLReader::FromFile(configFile);
    RapidXMLReaderNodePtr configNode = reader->getRoot("ConverterConfig");

    std::string type = configNode->attribute_value(XML::ATTRIBUTE_TYPE);
    if (type != NAME) throw Exception::XMLFormatException("ConverterConfig type: Expecting " + std::string(NAME) + ", got " + type);

    // Read Model Node
    RapidXMLReaderNodePtr modelNode = configNode->first_node("Model");
    if (outputModelProcessor) targetModelSize = outputModel->getHeight();
    else if (modelNode->has_attribute("size")) targetModelSize = XML::readFloat(modelNode->attribute_value("size").c_str());

    // Read Marker Mapping Node
    RapidXMLReaderNodePtr markerMappingNode = modelNode->first_node("MarkerMapping");
    for (RapidXMLReaderNodePtr mappingNode : markerMappingNode->nodes("Mapping")) {
        std::string inputMarkerName;
        if (mappingNode->has_attribute("c3d")) inputMarkerName = mappingNode->attribute_value("c3d"); // old term
        else inputMarkerName = mappingNode->attribute_value("source");
        std::string outputMarkerName;
        if (mappingNode->has_attribute("mmm")) outputMarkerName = mappingNode->attribute_value("mmm"); // old term
        else outputMarkerName = mappingNode->attribute_value("target");
        if (inputMarkerName.empty() || outputMarkerName.empty()) throw Exception::XMLFormatException("Empty Marker names not allowed!");
        for (const auto &mapping : markerMapping) {
            if (mapping.first == inputMarkerName) throw Exception::XMLFormatException("Mapping: Input marker name " + inputMarkerName + " is mapped twice!");
            else if (mapping.second == outputMarkerName) throw Exception::XMLFormatException("Mapping: Output marker name " + outputMarkerName + " is mapped twice!");
        }
        if (!outputModel->hasMarker(outputMarkerName)) throw Exception::XMLFormatException("Mapping: Output marker name " + outputMarkerName + " is invalid!");
        markerMapping[inputMarkerName] = outputMarkerName;
        float markerWeight = 1.0f;
        if (mappingNode->has_attribute("weight")) markerWeight = boost::lexical_cast<float>(mappingNode->attribute_value("weight"));
        markerWeights[inputMarkerName] = markerWeight;
    }
    // TODO check inputMarkerName in inputMotion/Model?

    // Read Joint Set Node
    if (modelNode->has_node("JointSet")) {
        RapidXMLReaderNodePtr jointSetNode = modelNode->first_node("JointSet");
        for (RapidXMLReaderNodePtr jointNode : jointSetNode->nodes("Joint")) {
            std::string jointName = jointNode->attribute_value(XML::ATTRIBUTE_NAME);
            if (!outputModel->getModelNode(jointName))
                throw Exception::XMLFormatException("Specified joint " + jointName + " in the converter configuration file is not part of the output model!");
            joints.push_back(jointName);
        }
    }

    // TODO OptimizationNode
}

void NloptMotionConverter::cancel() {
    if (strategy) strategy->cancel();
}

float NloptMotionConverter::getCurrentTimestep() {
    return strategy ? strategy->getCurrentTimestep() : 0;
}

MotionPtr NloptMotionConverter::convertMotion() {
    MMM_INFO << "Converting motion" << std::endl;

    // Create resulting output motion
    MotionPtr outputMotion(new Motion(convertedMotionName, outputModelUnprocessed, outputModel, outputModelProcessor));

    std::map<float, std::map<std::string, Eigen::Vector3f> > labeledMarkerData = getLabeledMarkerData();
//    std::map<float, std::map<std::string, Eigen::Vector3f> > _10framesData;

//    for(auto & d : labeledMarkerData)
//    {
//        _10framesData[d.first] = d.second;
//        if(_10framesData.size() >= 30)
//            break;
//    }
//    labeledMarkerData = _10framesData;
    // Use concrete Strategy
    switch (optimizationStrategy) {
    case FRAMEWISE_LOCAL:
        strategy = FramewiseLocalStrategyPtr(new FramewiseLocalStrategy(labeledMarkerData, outputMotion, outputModel, joints, markerMapping, markerWeights));
        break;
    case WHOLE_MOTION:
        MotionPtr localOutputMotion(new Motion(convertedMotionName, outputModelUnprocessed, outputModel, outputModelProcessor));

        strategy = FramewiseLocalStrategyPtr(new FramewiseLocalStrategy(labeledMarkerData, localOutputMotion, outputModel, joints, markerMapping, markerWeights));
        strategy->convert();

        auto globalStrategy = MinimumJerkStrategyPtr(new MinimumJerkStrategy(labeledMarkerData, outputMotion, outputModel, joints, markerMapping, markerWeights));
        globalStrategy->setInputMotion(localOutputMotion);
        strategy = globalStrategy;
    }
    strategy->convert();

    // Add other Sensors
    for (const auto &s : inputMotion->getSensorData()) {
        SensorPtr sensor = s.second;
        if (sensor->getType() != ModelPoseSensor::TYPE && sensor->getType() != KinematicSensor::TYPE) {
            outputMotion->addSensor(sensor); // exception should not occure!
        }
    }

    return outputMotion;
}


std::map<float, std::map<std::string, Eigen::Vector3f> > NloptMotionConverter::getLabeledMarkerData() {
    std::map<float, std::map<std::string, Eigen::Vector3f> > labeledMarkerData;
    ModelPtr inputModel = inputMotion->getModel();
    if (inputModel && inputMotion->hasSensor(ModelPoseSensor::TYPE) && inputMotion->hasSensor(KinematicSensor::TYPE)) {
        MMM_INFO << "Using motion's kinematic data." << std::endl;

        VirtualRobot::RobotPtr inputRobot;
        inputRobot = SimoxTools::buildModel(inputModel, false);
        if (!inputRobot) throw Exception::MMMException("Could not build Simox robot for input model!");
        inputRobot->setUpdateCollisionModel(false);
        inputRobot->setUpdateVisualization(false);
        inputRobot->setupVisualization(false, false);
        inputRobot->setThreadsafe(false);

        markerScaleFactor = targetModelSize / inputModel->getHeight();
        MMM_INFO << "Marker scale factor is " << markerScaleFactor << "." << std::endl;

        // Retrieve labeled marker data from modelPose and joints
        ModelPoseSensorPtr modelPoseSensor = inputMotion->getSensorByType<ModelPoseSensor>(ModelPoseSensor::TYPE);
        std::vector<float> timesteps = modelPoseSensor->getTimesteps();

        KinematicSensorPtr kinematicSensor = KinematicSensor::join(inputMotion->getSensorsByType<KinematicSensor>(KinematicSensor::TYPE), timesteps);
        const std::vector<std::string>& jointNames = kinematicSensor->getJointNames();

        for (auto timestep : timesteps)
        {
            ModelPoseSensorMeasurementPtr modelPoseSensorMeasurement = modelPoseSensor->getDerivedMeasurement(timestep);
            KinematicSensorMeasurementPtr kinematicSensorMeasurement = kinematicSensor->getDerivedMeasurement(timestep);
            if (!kinematicSensorMeasurement) continue;

            inputRobot->setGlobalPose(modelPoseSensorMeasurement->getRootPose());

            for (unsigned int i = 0; i < jointNames.size(); ++i) {
                inputRobot->setJointValue(jointNames[i], kinematicSensorMeasurement->getJointAngles()[i]);
            }

            std::map<std::string, Eigen::Vector3f> labeledMarker;

            for (const auto &mapping : markerMapping) {
                const std::string& inputMarker = mapping.first;
                VirtualRobot::SensorPtr sensor = inputRobot->getSensor(inputMarker);
                if (sensor) labeledMarker[inputMarker] = sensor->getGlobalPose().block(0, 3, 3, 1) * markerScaleFactor;
                else throw Exception::MMMException("The robot for the input motion doesn't contain the following input marker '" + inputMarker + "'! Maybe the wrong configuration file was choosen.");
            }

            labeledMarkerData[timestep] = labeledMarker;
        }
    } else if (inputMotion->hasSensor(MoCapMarkerSensor::TYPE)){
        MMM_INFO << "Using motion's marker data." << std::endl;

        MoCapMarkerSensorPtr moCapMarkerSensor = inputMotion->getSensorByType<MoCapMarkerSensor>(MoCapMarkerSensor::TYPE);
        labeledMarkerData = moCapMarkerSensor->getLabeledMarkerData();

        int missingMarker = 0;
        for (const auto &markerData : labeledMarkerData) {
            std::map<std::string, Eigen::Vector3f> m = markerData.second;
            for (auto it = markerMapping.begin(); it != markerMapping.end(); it++) {
                if (m.find(it->first) == m.end()) {
                    MMM_INFO << "Ignoring motion capture marker "<< it->first << ", because isn't contained in the motion (timestep: " << markerData.first << ")!" << std::endl;
                    markerMapping.erase(it);
                    missingMarker++;
                }
            }
        }
        if (missingMarker > 3) throw Exception::MMMException("Too much motion capture marker are missing in the motion!");
    }
    else throw Exception::MMMException("Either MoCapMarkerSensor or a Model + ModelPoseSensor + KinematicSensor is needed!");
    return labeledMarkerData;
}

